{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G3MMAcssHTML"
      },
      "source": [
        "<link rel=\"stylesheet\" href=\"/site-assets/css/gemma.css\">\n",
        "<link rel=\"stylesheet\" href=\"https://fonts.googleapis.com/css2?family=Google+Symbols:opsz,wght,FILL,GRAD@20..48,100..700,0..1,-50..200\" />"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2024 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SDEExiAk4fLb"
      },
      "source": [
        "# AI Assisted programming with CodeGemma and KerasNLP"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZFWzQEqNosrS"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/gemma/docs/codegemma/code_assist_keras\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on ai.google.dev</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/codegemma/code_assist_keras.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/generative-ai-docs/blob/main/site/en/gemma/docs/codegemma/code_assist_keras.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lSGRSsRPgkzK"
      },
      "source": [
        "## Overview\n",
        "\n",
        "CodeGemma is a variant of Gemma that is fine-tuned for coding tasks. This tutorial builds on the [Keras CodeGemma quickstart](https://colab.research.google.com/drive/11Va7W2Yl12JUnfx9YDJmy90kBu1i4rtw) and shows you more ways in which CodeGemma can assist your programming tasks.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AraV8lsDaP1A"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lyhHCMfoRZ_v"
      },
      "source": [
        "### Get access to CodeGemma\n",
        "\n",
        "To complete this tutorial, you will first need to complete the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup). The Gemma setup instructions show you how to do the following:\n",
        "\n",
        "* Get access to Gemma on [kaggle.com](https://kaggle.com){:.external}.\n",
        "* Select a Colab runtime with sufficient resources to run the Gemma 7B model.\n",
        "* Generate and configure a Kaggle username and API key.\n",
        "\n",
        "After you've completed the Gemma setup, move on to the next section, where you'll set environment variables for your Colab environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AZ5Qo0fxRZ1V"
      },
      "source": [
        "### Select the runtime\n",
        "\n",
        "To run the CodeGemma 7B models, you'll need to have a paid Colab Pro plan which provides a runtime with an A100 GPU.\n",
        "\n",
        "1. In the upper-right of the Colab window, select &#9662; (**Additional connection options**).\n",
        "2. Select **Change runtime type**.\n",
        "3. Under **Hardware accelerator**, select **A100 GPU**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hsPC0HRkJl0K"
      },
      "source": [
        "### Configure your API key\n",
        "\n",
        "To use Gemma, you must provide your Kaggle username and a Kaggle API key.\n",
        "\n",
        "To generate a Kaggle API key, go to the **Account** tab of your Kaggle user profile and select **Create New Token**. This will trigger the download of a `kaggle.json` file containing your API credentials.\n",
        "\n",
        "In Colab, select **Secrets** (🔑) in the left pane and add your Kaggle username and Kaggle API key. Store your username under the name `KAGGLE_USERNAME` and your API key under the name `KAGGLE_KEY`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7iOF6Yo-wUEC"
      },
      "source": [
        "### Set environment variables\n",
        "\n",
        "Set environment variables for `KAGGLE_USERNAME` and `KAGGLE_KEY`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DrBoa_Urw9Vx"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "\n",
        "os.environ[\"KAGGLE_USERNAME\"] = userdata.get('KAGGLE_USERNAME')\n",
        "os.environ[\"KAGGLE_KEY\"] = userdata.get('KAGGLE_KEY')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FX47AUYrXwLK"
      },
      "source": [
        "### Install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ACnm31nfBqHj"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m508.4/508.4 kB\u001b[0m \u001b[31m3.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m950.8/950.8 kB\u001b[0m \u001b[31m18.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.2/5.2 MB\u001b[0m \u001b[31m51.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m589.8/589.8 MB\u001b[0m \u001b[31m1.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.8/4.8 MB\u001b[0m \u001b[31m59.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.2/2.2 MB\u001b[0m \u001b[31m32.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.5/5.5 MB\u001b[0m \u001b[31m42.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m45.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m311.2/311.2 kB\u001b[0m \u001b[31m38.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "tf-keras 2.15.1 requires tensorflow<2.16,>=2.15, but you have tensorflow 2.16.1 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "!pip install -q -U keras-nlp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2I69cArSBm3z"
      },
      "source": [
        "### Select a backend"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QPTF92kyOQ-p"
      },
      "source": [
        "Keras is a high-level, multi-framework deep learning API designed for simplicity and ease of use. Using Keras 3, you can run workflows on one of three backends: TensorFlow, JAX, or PyTorch.\n",
        "\n",
        "For this tutorial, configure the backend for JAX."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ww83zI9ToPso"
      },
      "outputs": [],
      "source": [
        "os.environ[\"KERAS_BACKEND\"] = \"jax\"  # Or \"tensorflow\" or \"torch\"."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FLDJd1nxa3I7"
      },
      "source": [
        "### Import packages\n",
        "\n",
        "Import Keras and KerasNLP."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oQkqsyE1a2YD"
      },
      "outputs": [],
      "source": [
        "import keras_nlp\n",
        "import keras\n",
        "\n",
        "# Run at half precision.\n",
        "keras.config.set_floatx(\"bfloat16\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dn_yBThs4kXk"
      },
      "source": [
        "## CodeGemma 7B Model Examples\n",
        "\n",
        "This section covers examples of using the pre-trained 7B CodeGemma model to help with coding tasks.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7RCE3fdGhDE5"
      },
      "source": [
        "### Load the model\n",
        "\n",
        "KerasNLP provides implementations of all three CodeGemma variants (2B and 7B pre-trained (PT) and 7B instruction-tuned (IT)) using [`GemmaCausalLM`](https://keras.io/api/keras_nlp/models/gemma/gemma_causal_lm/){:.external}, an end-to-end Gemma model for causal language modeling. A causal language model predicts the next token based on previous tokens."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S43GTKpQF_-F"
      },
      "source": [
        "For this example, load the `code_gemma_7b_en` model using the [`from_preset`](https://keras.io/api/keras_nlp/models/gemma/gemma_causal_lm/#frompreset-method){:.external} method.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SQ6Yr7c4GKTx"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/config.json...\n",
            "100%|██████████| 556/556 [00:00<00:00, 790kB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/model.weights.h5...\n",
            "100%|██████████| 15.9G/15.9G [02:39<00:00, 107MB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/tokenizer.json...\n",
            "100%|██████████| 401/401 [00:00<00:00, 587kB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_7b_en/1/download/assets/tokenizer/vocabulary.spm...\n",
            "100%|██████████| 4.04M/4.04M [00:00<00:00, 16.4MB/s]\n"
          ]
        }
      ],
      "source": [
        "gemma_lm_7b = keras_nlp.models.GemmaCausalLM.from_preset(\"code_gemma_7b_en\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v0YFsh2a3n9P"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n",
              "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ],
            "text/plain": [
              "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n",
              "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">3072</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">8,537,680,896</span> │ padding_mask[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>],        │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaBackbone</span>)               │                           │                 │ token_ids[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]            │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_embedding               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">256000</span>)      │     <span style=\"color: #00af00; text-decoration-color: #00af00\">786,432,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n",
              "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n",
              "</pre>\n"
            ],
            "text/plain": [
              "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m3072\u001b[0m)        │   \u001b[38;5;34m8,537,680,896\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m],        │\n",
              "│ (\u001b[38;5;33mGemmaBackbone\u001b[0m)               │                           │                 │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]            │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_embedding               │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m)      │     \u001b[38;5;34m786,432,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n",
              "│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n",
              "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">8,537,680,896</span> (15.90 GB)\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m8,537,680,896\u001b[0m (15.90 GB)\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">8,537,680,896</span> (15.90 GB)\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m8,537,680,896\u001b[0m (15.90 GB)\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> (0.00 B)\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m0\u001b[0m (0.00 B)\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "gemma_lm_7b.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YXO1l87PtfSN"
      },
      "source": [
        "The `from_preset` method instantiates the model from a preset architecture and weights."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H2XNKl2vHxHy"
      },
      "source": [
        "### Code completion with Multi-line FIM\n",
        "\n",
        "The PT CodeGemma models are trained on code infilling tasks. This section shows examples that use CodeGemma's multi-line fill-in the-middle (FIM) capability to autofill code at the specified cursor location based on the surrounding context.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Gsjz0QwSz2f"
      },
      "source": [
        "As a first step, define constants and a prompt formatting helper function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mHAjAdvB_5yN"
      },
      "outputs": [],
      "source": [
        "# Formatting control tokens to specify cursor location\n",
        "BEFORE_CURSOR = \"<|fim_prefix|>\"\n",
        "AFTER_CURSOR = \"<|fim_suffix|>\"\n",
        "AT_CURSOR = \"<|fim_middle|>\"\n",
        "FILE_SEPARATOR = \"<|file_separator|>\"\n",
        "\n",
        "# Define model stop tokens\n",
        "END_TOKEN = gemma_lm_7b.preprocessor.tokenizer.end_token\n",
        "stop_tokens = (BEFORE_CURSOR, AFTER_CURSOR, AT_CURSOR, FILE_SEPARATOR, END_TOKEN)\n",
        "stop_token_ids = tuple(gemma_lm_7b.preprocessor.tokenizer.token_to_id(x) for x in stop_tokens)\n",
        "\n",
        "def format_completion_prompt(before, after):\n",
        "    return f\"{BEFORE_CURSOR}{before}{AFTER_CURSOR}{after}{AT_CURSOR}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aDwFjBhoNe02"
      },
      "source": [
        "#### Example 1 - Insert missing condition\n",
        "\n",
        "The example code below to generate the Fibonacci sequence will not execute correctly if `n=1`:\n",
        "\n",
        "```python\n",
        "def fibonacci(n: int) -> int:\n",
        "  if n == 0:\n",
        "    return 0\n",
        "  # The cursor is right before the e in the following line\n",
        "  else:\n",
        "    return fibonacci(n - 1) + fibonacci(n - 2)\n",
        "```\n",
        "\n",
        "Assuming that the cursor is at the beginning of line 4 (where the `else` clause is), then the content before and after the cursor is:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eGU7XXk1SFYT"
      },
      "outputs": [],
      "source": [
        "before = \"\"\"def fibonacci(n: int) -> int:\\n  if n == 0:\\n    return 0\\n\"\"\" # Mind the spaces!\n",
        "after = \"\"\"\\n  else:\\n    return fibonacci(n - 1) + fibonacci(n-2)\\n\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GoRC_SzuAO7t"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<|fim_prefix|>def fibonacci(n: int) -> int:\n",
            "  if n == 0:\n",
            "    return 0\n",
            "<|fim_suffix|>\n",
            "  else:\n",
            "    return fibonacci(n - 1) + fibonacci(n-2)\n",
            "<|fim_middle|>\n"
          ]
        }
      ],
      "source": [
        "prompt = format_completion_prompt(before, after)\n",
        "print(prompt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JRI9npXOAO7t"
      },
      "source": [
        "Run the prompt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MpRJDxf4AO7t"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<|fim_prefix|>def fibonacci(n: int) -> int:\n",
            "  if n == 0:\n",
            "    return 0\n",
            "<|fim_suffix|>\n",
            "  else:\n",
            "    return fibonacci(n - 1) + fibonacci(n-2)\n",
            "<|fim_middle|>elif n == 1:\n",
            "    return 1<|file_separator|>\n"
          ]
        }
      ],
      "source": [
        "print(gemma_lm_7b.generate(prompt, stop_token_ids=stop_token_ids, max_length=128))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "np8zhxowfEsB"
      },
      "source": [
        "The model inserts the correct `elif` conidtion for `n=1` at the location of the cursor."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AliD_nTEeBMV"
      },
      "source": [
        "#### Example 2 - Complete DFS traversal algorithm\n",
        "\n",
        "Auto-complete code for a depth-first search (DFS) tree traversal algorithm."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kjaCaWcxqoq5"
      },
      "outputs": [],
      "source": [
        "before = \"\"\"void dfs(node* root) {\n",
        "  if (root->left) {\n",
        "    dfs(root->left);\n",
        "  }\"\"\"\n",
        "after = \"\"\"\\nprintf(\"%d\", root->value);\n",
        "}\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nq8uV7DehsOt"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<|fim_prefix|>void dfs(node* root) {\n",
            "  if (root->left) {\n",
            "    dfs(root->left);\n",
            "  }<|fim_suffix|>\n",
            "printf(\"%d\", root->value);\n",
            "}<|fim_middle|>\n"
          ]
        }
      ],
      "source": [
        "prompt = format_completion_prompt(before, after)\n",
        "print(prompt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wexpoZS9hsOu"
      },
      "source": [
        "Run the prompt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mCY-VI01hsOu"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<|fim_prefix|>void dfs(node* root) {\n",
            "  if (root->left) {\n",
            "    dfs(root->left);\n",
            "  }<|fim_suffix|>\n",
            "printf(\"%d\", root->value);\n",
            "}<|fim_middle|>\n",
            "  if (root->right) {\n",
            "    dfs(root->right);\n",
            "  }<|file_separator|>\n"
          ]
        }
      ],
      "source": [
        "print(gemma_lm_7b.generate(prompt, stop_token_ids=stop_token_ids, max_length=128))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v7l8rVt0-irP"
      },
      "source": [
        "### Code generation\n",
        "\n",
        "In addition to code infilling, the CodeGemma 7B PT is model is also trained on natural language corpuses. You can use this to prompt the model to generate code.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ECfsJ7JW-7ef"
      },
      "outputs": [],
      "source": [
        "generation_prompt= \"\"\"Write a rust function to identify non-prime numbers.\n",
        "Examples:\n",
        ">>> is_not_prime(2)\n",
        "False\n",
        ">>> is_not_prime(10)\n",
        "True\n",
        "pub fn is_not_prime(n: i32) -> bool {\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-RZkx_0i_Vgj"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Write a rust function to identify non-prime numbers.\n",
            "Examples:\n",
            ">>> is_not_prime(2)\n",
            "False\n",
            ">>> is_not_prime(10)\n",
            "True\n",
            "pub fn is_not_prime(n: i32) -> bool {\n",
            "    if n <= 1 {\n",
            "        return true;\n",
            "    }\n",
            "    for i in 2..n {\n",
            "        if n % i == 0 {\n",
            "            return true;\n",
            "        }\n",
            "    }\n",
            "    false\n",
            "}\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(gemma_lm_7b.generate(generation_prompt, max_length=500))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tzs6i4cieb2Y"
      },
      "source": [
        "## 7B IT model examples\n",
        "\n",
        "This section uses the CodeGemma 7B Instruction-Tuned model for more advanced coding tasks. The CodeGemma 7B IT model is derived from the CodeGemma 7B PT model through supervised fine-tuning on code along with Reinforcement Learning with Human Feedback. This section covers examples of using this model for open-ended generation.\n",
        "\n",
        "NOTE: If you are working through this tutorial in Colab and already have the 2B model loaded from above, restart your Colab Runtime by going to **Runtime** > **Disconnect and delete runtime** and re-connect to a new runtime. This frees up memory and prevents out-of-memory (OOM) issues. After connecting to a new runtime, re-run the Setup steps [from here](#set_environment_variables) before proceeding further.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wsNKfK4iGYRR"
      },
      "source": [
        "### Load the IT model\n",
        "\n",
        "Load the `code_gemma_instruct_7b_en` model using the `from_preset` method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HxF9z5cFG5Yz"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/config.json...\n",
            "100%|██████████| 556/556 [00:00<00:00, 754kB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/model.weights.h5...\n",
            "100%|██████████| 15.9G/15.9G [03:18<00:00, 86.2MB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/tokenizer.json...\n",
            "100%|██████████| 401/401 [00:00<00:00, 593kB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/codegemma/keras/code_gemma_instruct_7b_en/1/download/assets/tokenizer/vocabulary.spm...\n",
            "100%|██████████| 4.04M/4.04M [00:00<00:00, 16.8MB/s]\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n",
              "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n",
              "</pre>\n"
            ],
            "text/plain": [
              "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n",
              "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">3072</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">8,537,680,896</span> │ padding_mask[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>],        │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaBackbone</span>)               │                           │                 │ token_ids[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]            │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_embedding               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">256000</span>)      │     <span style=\"color: #00af00; text-decoration-color: #00af00\">786,432,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n",
              "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n",
              "</pre>\n"
            ],
            "text/plain": [
              "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m3072\u001b[0m)        │   \u001b[38;5;34m8,537,680,896\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m],        │\n",
              "│ (\u001b[38;5;33mGemmaBackbone\u001b[0m)               │                           │                 │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]            │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_embedding               │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m)      │     \u001b[38;5;34m786,432,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n",
              "│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n",
              "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">8,537,680,896</span> (15.90 GB)\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m8,537,680,896\u001b[0m (15.90 GB)\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">8,537,680,896</span> (15.90 GB)\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m8,537,680,896\u001b[0m (15.90 GB)\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> (0.00 B)\n",
              "</pre>\n"
            ],
            "text/plain": [
              "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m0\u001b[0m (0.00 B)\n"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "gemma_lm_7b_it = keras_nlp.models.GemmaCausalLM.from_preset(\"code_gemma_instruct_7b_en\")\n",
        "gemma_lm_7b_it.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x-KDJ0KvcM9S"
      },
      "source": [
        "IT models are trained with a specific formatter that annotates all instruction tuning examples with extra information to indicate roles and delineate turns in a conversation.\n",
        "\n",
        "As a first step, define constants and a prompt formatting helper function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "llseiN-_5XXk"
      },
      "outputs": [],
      "source": [
        "# Formatting control tokens for instruction tuning\n",
        "START_OF_TURN_USER = \"<start_of_turn>user\"\n",
        "END_OF_TURN = \"<end_of_turn>\"\n",
        "START_OF_TURN_MODEL = \"<start_of_turn>model\"\n",
        "\n",
        "# Formatting helper function\n",
        "def format_instruction_prompt(context):\n",
        "    return f\"{START_OF_TURN_USER}\\n{context}{END_OF_TURN}\\n{START_OF_TURN_MODEL}\\n\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m9H0F8k5a8OO"
      },
      "source": [
        "### Code translation\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-RJ6z5gbqTPA"
      },
      "outputs": [],
      "source": [
        "context1 = \"\"\"\n",
        "You are an experienced C and Python programmer. Convert the following Python code into C.\n",
        "```python\n",
        "def factorial(n):\n",
        "    result = 1\n",
        "    for i in range(2, n + 1):\n",
        "        result *= i\n",
        "    return result\n",
        "```\\n\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eg2TcYZFYfiU"
      },
      "source": [
        "Format the prompt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LlDJmIAz5f0R"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<start_of_turn>user\n",
            "\n",
            "You are an experienced C and Python programmer. Convert the following Python code into C.\n",
            "```python\n",
            "def factorial(n):\n",
            "    result = 1\n",
            "    for i in range(2, n + 1):\n",
            "        result *= i\n",
            "    return result\n",
            "```\n",
            "<end_of_turn>\n",
            "<start_of_turn>model\n",
            "\n"
          ]
        }
      ],
      "source": [
        "prompt1 = format_instruction_prompt(context1)\n",
        "print(prompt1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xPE0-SGP68dP"
      },
      "source": [
        "Run the prompt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rSBsedwN676V"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<start_of_turn>user\n",
            "\n",
            "You are an experienced C and Python programmer. Convert the following Python code into C.\n",
            "```python\n",
            "def factorial(n):\n",
            "    result = 1\n",
            "    for i in range(2, n + 1):\n",
            "        result *= i\n",
            "    return result\n",
            "```\n",
            "<end_of_turn>\n",
            "<start_of_turn>model\n",
            "Here is the C code equivalent of the Python code:\n",
            "\n",
            "```c\n",
            "int factorial(int n) {\n",
            "  int result = 1;\n",
            "  for (int i = 2; i <= n; i++) {\n",
            "    result *= i;\n",
            "  }\n",
            "  return result;\n",
            "}\n",
            "```\n",
            "\n",
            "Here is a breakdown of the changes:\n",
            "\n",
            "* The function is declared with the `int` return type, as in Python.\n",
            "* The `for` loop is converted to a `for` loop with an `int` variable `i` initialized to 2 and incremented by 1 in each iteration.\n",
            "* The `range` function is replaced with a simple loop that iterates from 2 to `n` (inclusive).\n",
            "* The `result *= i` statement is used to multiply `result` by `i` in each iteration.\n",
            "* The `return` statement is used to return the final value of `result`.\n"
          ]
        }
      ],
      "source": [
        "print(gemma_lm_7b_it.generate(prompt1, max_length=500))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0b8RXp7TPpEi"
      },
      "source": [
        "### Code vulnerability detection\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yYW6VXt7qPTc"
      },
      "outputs": [],
      "source": [
        "context2 = \"\"\"\n",
        "You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.\n",
        "```cpp\n",
        "int i;\n",
        "unsigned int numWidgets;\n",
        "Widget **WidgetList;\n",
        "\n",
        "numWidgets = GetUntrustedSizeValue();\n",
        "if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {\n",
        "    ExitError(\"Incorrect number of widgets requested!\");\n",
        "}\n",
        "WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));\n",
        "printf(\"WidgetList ptr=%p\\n\", WidgetList);\n",
        "for (i = 0; i < numWidgets; i++) {\n",
        "    WidgetList[i] = InitializeWidget();\n",
        "}\n",
        "WidgetList[numWidgets] = NULL;\n",
        "showWidgets(WidgetList);\n",
        "```\\n\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "crZ2DWaczme4"
      },
      "source": [
        "Format the prompt."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "To_5KvKJPw1H"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<start_of_turn>user\n",
            "\n",
            "You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.\n",
            "```cpp\n",
            "int i;\n",
            "unsigned int numWidgets;\n",
            "Widget **WidgetList;\n",
            "\n",
            "numWidgets = GetUntrustedSizeValue();\n",
            "if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {\n",
            "    ExitError(\"Incorrect number of widgets requested!\");\n",
            "}\n",
            "WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));\n",
            "printf(\"WidgetList ptr=%p\n",
            "\", WidgetList);\n",
            "for (i = 0; i < numWidgets; i++) {\n",
            "    WidgetList[i] = InitializeWidget();\n",
            "}\n",
            "WidgetList[numWidgets] = NULL;\n",
            "showWidgets(WidgetList);\n",
            "```\n",
            "<end_of_turn>\n",
            "<start_of_turn>model\n",
            "\n"
          ]
        }
      ],
      "source": [
        "prompt2 = format_instruction_prompt(context2)\n",
        "print(prompt2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jHZg8YX8QuIe"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<start_of_turn>user\n",
            "\n",
            "You are an experienced C++ programmer hunting for vulnerable code. Is the following code vulnerable? Explain your reasoning.\n",
            "```cpp\n",
            "int i;\n",
            "unsigned int numWidgets;\n",
            "Widget **WidgetList;\n",
            "\n",
            "numWidgets = GetUntrustedSizeValue();\n",
            "if ((numWidgets == 0) || (numWidgets > MAX_NUM_WIDGETS)) {\n",
            "    ExitError(\"Incorrect number of widgets requested!\");\n",
            "}\n",
            "WidgetList = (Widget **) malloc(numWidgets * sizeof(Widget *));\n",
            "printf(\"WidgetList ptr=%p\n",
            "\", WidgetList);\n",
            "for (i = 0; i < numWidgets; i++) {\n",
            "    WidgetList[i] = InitializeWidget();\n",
            "}\n",
            "WidgetList[numWidgets] = NULL;\n",
            "showWidgets(WidgetList);\n",
            "```\n",
            "<end_of_turn>\n",
            "<start_of_turn>model\n",
            "Yes, the code is vulnerable to a memory access error.\n",
            "\n",
            "**Reasoning:**\n",
            "\n",
            "* The code allocates memory for `WidgetList` using `malloc` based on the value of `numWidgets`.\n",
            "* However, the loop iterates from `0` to `numWidgets`, which is one element beyond the allocated memory.\n",
            "* This means that accessing `WidgetList[numWidgets]` will result in a memory access error, as it is outside the bounds of the allocated memory.\n",
            "\n",
            "**Example of Memory Access Error:**\n",
            "\n",
            "When `numWidgets` is 5, the code allocates memory for `WidgetList` as follows:\n",
            "\n",
            "```\n",
            "WidgetList = (Widget **) malloc(5 * sizeof(Widget *));\n",
            "```\n",
            "\n",
            "The loop iterates from 0 to 4, accessing the following elements:\n",
            "\n",
            "* `WidgetList[0]`\n",
            "* `WidgetList[1]`\n",
            "* `WidgetList[2]`\n",
            "* `WidgetList[3]`\n",
            "* `WidgetList[4]`\n",
            "\n",
            "However, the code then attempts to access `WidgetList[5]`, which is outside the allocated memory range. This will result in a memory access error.\n",
            "\n",
            "**Solution:**\n",
            "\n",
            "To resolve this vulnerability, the loop should be modified to iterate from 0 to `numWidgets - 1`:\n",
            "\n",
            "```cpp\n",
            "for (i = 0; i < numWidgets - 1; i++) {\n",
            "    WidgetList[i] = InitializeWidget();\n",
            "}\n",
            "```\n",
            "\n",
            "This ensures that the loop does not access elements beyond the allocated memory range.\n"
          ]
        }
      ],
      "source": [
        "print(gemma_lm_7b_it.generate(prompt2, max_length=1000))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GRVutCnnew2N"
      },
      "source": [
        "The model detects a potential vulnerability in the code and provides code changes to mitigate it."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QesAFYNcRw5f"
      },
      "source": [
        "## Summary\n",
        "\n",
        "This tutorial walked you through using CodeGemma for a variety of coding tasks. To learn more about CodeGemma:\n",
        "\n",
        "* Refer to the [CodeGemma model card](https://ai.google.dev/gemma/docs/codegemma/model_card) for the technical specs of the CodeGemma models.\n",
        "* Learn more about how to use CodeGemma in VertexAI [here](https://colab.sandbox.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb).\n",
        "* Check out the [Keras CodeGemma quickstart](https://ai.google.dev/gemma/docs/codegemma/keras_quickstart)."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "code_assist_keras.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
